from typing import Dict, Any, List
from .base_agent import BaseAgent
from utils import logger, CodeAnalyzer
import json

class OptimizerAgent(BaseAgent):
    """优化专家代理 - 负责性能优化和代码改进"""
    
    def __init__(self):
        super().__init__(
            name="优化专家",
            role="性能优化专家和代码改进师"
        )
        self.code_analyzer = CodeAnalyzer()
    
    def _get_system_prompt(self) -> str:
        return """你是一个资深的性能优化专家，具有以下专业技能：

⚡ **性能优化**：
- 算法复杂度优化
- 内存使用优化
- 数据库查询优化
- 缓存策略设计
- 并发和异步处理优化

🔧 **代码改进**：
- 代码重构和清理
- 设计模式应用
- 可维护性提升
- 可读性改进
- 错误处理增强

📊 **性能分析**：
- 性能瓶颈识别
- 资源使用分析
- 负载测试和压力测试
- 监控和指标设计
- 性能基准建立

🛠️ **优化技术**：
- 数据结构选择优化
- 算法替换和改进
- 多线程和多进程优化
- I/O操作优化
- 网络通信优化

🎯 **优化目标**：
- 响应时间改进
- 吞吐量提升
- 资源消耗降低
- 可扩展性增强
- 用户体验提升

请提供实用、有效的优化方案和改进建议。用中文回复。"""
    
    def process(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """分析代码并提供优化建议"""
        generated_code = input_data.get('generated_code', '')
        code_files = input_data.get('code_files', {})
        review_report = input_data.get('review_report', '')
        test_results = input_data.get('test_results', {})
        
        if not generated_code and not code_files:
            return {
                'success': False,
                'error': '未提供待优化的代码'
            }
        
        # 如果没有解析的代码文件，尝试从生成的代码中解析
        if not code_files and generated_code:
            code_files = {'main.py': generated_code}
        
        # 分析性能瓶颈
        performance_analysis = self._analyze_performance(code_files)
        
        # 构造优化提示
        prompt = f"""
基于以下信息进行性能分析和代码优化：

**代码文件**：
{self._format_code_files(code_files)}

**性能分析**：
{performance_analysis}

**审查报告摘要**：
{review_report[:500] + '...' if len(review_report) > 500 else review_report}

**测试结果**：
{str(test_results) if test_results else '暂无测试结果'}

**请提供全面的优化方案，包括**：

## 1. 性能瓶颈分析
- 关键性能问题识别
- 瓶颈根本原因分析
- 性能影响评估
- 优化优先级排序

## 2. 算法优化
- 时间复杂度改进方案
- 空间复杂度优化建议
- 数据结构选择优化
- 算法替换建议

## 3. 代码优化
- 循环和条件优化
- 函数调用优化
- 内存分配优化
- I/O操作优化

## 4. 架构优化
- 模块结构改进
- 缓存策略设计
- 并发处理优化
- 数据库访问优化

## 5. 具体优化代码
为主要瓶颈提供优化后的代码实现，格式：

```python
# === 优化代码: 文件名_优化 ===
# 优化说明

优化后的代码...
```

## 6. 性能监控建议
- 关键指标定义
- 监控点设置
- 性能基准建立
- 警报和通知机制

请提供详细的优化方案和实施步骤。
"""
        
        try:
            optimization_plan = self._call_llm(prompt)
            optimized_code = self._extract_optimized_code(optimization_plan)
            
            # 生成性能改进摘要
            improvement_summary = self._generate_improvement_summary(performance_analysis, optimization_plan)
            
            return {
                'success': True,
                'optimization_plan': optimization_plan,
                'optimized_code': optimized_code,
                'performance_analysis': performance_analysis,
                'improvement_summary': improvement_summary,
                'summary': f'优化分析完成，识别出{len(performance_analysis.get("bottlenecks", []))}个性能瓶颈'
            }
            
        except Exception as e:
            logger.error(f"性能优化分析失败: {e}")
            return {
                'success': False,
                'error': str(e),
                'summary': '性能优化分析失败'
            }
    
    def _analyze_performance(self, code_files: Dict[str, str]) -> Dict[str, Any]:
        """分析代码性能特征"""
        try:
            # 使用模型分析性能
            analysis_prompt = f"""
你是一个性能分析专家。请分析以下代码文件的性能特征，并以JSON格式返回分析结果。

代码文件：
{json.dumps(code_files, indent=2, ensure_ascii=False)}

你必须严格按照以下JSON格式返回分析结果，不要添加任何其他内容：
{{
    "bottlenecks": [
        {{
            "file": "",
            "issue": "",
            "severity": "",
            "line_number": 0,
            "suggestion": ""
        }}
    ],
    "complexity_issues": [
        {{
            "file": "",
            "issue": "",
            "severity": "",
            "line_number": 0,
            "suggestion": ""
        }}
    ],
    "memory_issues": [
        {{
            "file": "",
            "issue": "",
            "severity": "",
            "line_number": 0,
            "suggestion": ""
        }}
    ],
    "io_issues": [
        {{
            "file": "",
            "issue": "",
            "severity": "",
            "line_number": 0,
            "suggestion": ""
        }}
    ],
    "optimization_opportunities": [
        {{
            "file": "",
            "type": "",
            "description": "",
            "priority": "",
            "suggestion": ""
        }}
    ]
}}

重要提示：
1. 你的回复必须是一个有效的JSON字符串
2. 不要添加任何解释、说明或其他文本
3. 所有数值必须是数字类型（不要用字符串）
4. 所有字符串必须用双引号（不要用单引号）
5. 不要使用任何注释
6. 不要使用任何换行符或空格（除了JSON格式必需的）
7. 确保JSON格式完全正确，可以被json.loads()解析
8. severity 和 priority 字段只能是 "high"、"medium" 或 "low"
"""
            
            analysis = self._call_llm(analysis_prompt)
            
            # 尝试清理响应中的非JSON内容
            analysis = analysis.strip()
            if analysis.startswith('```json'):
                analysis = analysis[7:]
            if analysis.endswith('```'):
                analysis = analysis[:-3]
            analysis = analysis.strip()
            
            try:
                return json.loads(analysis)
            except json.JSONDecodeError as e:
                logger.error(f"JSON解析失败: {e}")
                logger.error(f"原始响应: {analysis}")
                # 返回默认值
                return {
                    'error': f"JSON解析失败: {e}",
                    'bottlenecks': [],
                    'complexity_issues': [],
                    'memory_issues': [],
                    'io_issues': [],
                    'optimization_opportunities': []
                }
            
        except Exception as e:
            logger.error(f"性能分析失败: {e}")
            return {
                'bottlenecks': [],
                'complexity_issues': [],
                'memory_issues': [],
                'io_issues': [],
                'optimization_opportunities': []
            }
    
    def _detect_nested_loops(self, code_content: str) -> List[Dict[str, Any]]:
        """检测嵌套循环"""
        issues = []
        lines = code_content.split('\n')
        
        for i, line in enumerate(lines):
            if 'for ' in line and 'in ' in line:
                # 检查后续行是否有嵌套循环
                for j in range(i + 1, min(i + 20, len(lines))):
                    if '    for ' in lines[j] or '\tfor ' in lines[j]:
                        issues.append({
                            'type': 'nested_loop',
                            'line': i + 1,
                            'description': '检测到嵌套循环，可能影响性能'
                        })
                        break
        
        return issues
    
    def _detect_io_operations(self, code_content: str) -> List[Dict[str, Any]]:
        """检测I/O操作"""
        issues = []
        io_patterns = ['open(', 'file.read()', 'file.write()', 'requests.get(', 'requests.post(']
        
        lines = code_content.split('\n')
        for i, line in enumerate(lines):
            for pattern in io_patterns:
                if pattern in line:
                    issues.append({
                        'type': 'io_operation',
                        'line': i + 1,
                        'pattern': pattern,
                        'description': f'检测到I/O操作：{pattern}'
                    })
        
        return issues
    
    def _detect_memory_issues(self, code_content: str) -> List[Dict[str, Any]]:
        """检测内存相关问题"""
        issues = []
        memory_patterns = ['list(' * 100, 'dict(' * 100, '+ [' * 10]  # 简化的模式检测
        
        lines = code_content.split('\n')
        for i, line in enumerate(lines):
            # 检测大列表创建
            if 'range(' in line and any(str(n) in line for n in ['1000', '10000', '100000']):
                issues.append({
                    'type': 'large_range',
                    'line': i + 1,
                    'description': '检测到大范围创建，可能消耗过多内存'
                })
            
            # 检测列表拼接
            if '+= [' in line or '+ [' in line:
                issues.append({
                    'type': 'list_concatenation',
                    'line': i + 1,
                    'description': '检测到列表拼接，建议使用extend()或生成器'
                })
        
        return issues
    
    def _format_code_files(self, code_files: Dict[str, str]) -> str:
        """格式化代码文件为字符串"""
        formatted = ""
        for file_path, content in code_files.items():
            formatted += f"\n--- {file_path} ---\n"
            # 只显示前1000个字符以避免过长
            content_preview = content[:1000]
            if len(content) > 1000:
                content_preview += f"\n... [文件内容已截断，总长度: {len(content)} 字符]"
            formatted += content_preview + "\n"
        return formatted
    
    def _extract_optimized_code(self, optimization_plan: str) -> Dict[str, str]:
        """从优化计划中提取优化后的代码"""
        optimized_files = {}
        lines = optimization_plan.split('\n')
        current_file = None
        current_content = []
        
        for line in lines:
            if line.strip().startswith('# === 优化代码:') and line.strip().endswith('==='):
                # 保存上一个文件
                if current_file and current_content:
                    optimized_files[current_file] = '\n'.join(current_content)
                
                # 开始新文件
                current_file = line.strip().replace('# === 优化代码:', '').replace('===', '').strip()
                current_content = []
            elif current_file and line.startswith('```'):
                continue  # 跳过代码块结束标记
            elif current_file:
                current_content.append(line)
        
        # 保存最后一个文件
        if current_file and current_content:
            optimized_files[current_file] = '\n'.join(current_content)
        
        return optimized_files
    
    def _parse_optimized_files(self, response: str) -> Dict[str, str]:
        """解析优化后的代码文件"""
        optimized_files = {}
        current_file = None
        current_content = []
        
        for line in response.split('\n'):
            if line.startswith('# === 优化代码:'):
                # 保存之前的文件
                if current_file and current_content:
                    optimized_files[current_file] = '\n'.join(current_content)
                # 开始新文件
                current_file = line.strip().replace('# === 优化代码:', '').replace('===', '').strip()
                current_content = []
            elif current_file and line.startswith('```python'):
                continue  # 跳过代码块标记
            elif current_file and line.startswith('```'):
                continue  # 跳过代码块结束标记
            elif current_file:
                current_content.append(line)
        
        # 保存最后一个文件
        if current_file and current_content:
            optimized_files[current_file] = '\n'.join(current_content)
        
        return optimized_files

    def _generate_improvement_summary(self, performance_analysis: Dict[str, Any], optimization_plan: str) -> str:
        """生成性能改进摘要"""
        # 实现生成性能改进摘要的逻辑
        pass
